"""
Differential operators for grids.

This module provides both legacy operator functions and integration with
the new comprehensive operator system. For new code, use the operators
submodule which provides enhanced functionality.
"""

using LinearAlgebra
using SparseArrays
using FFTW

# Include comprehensive operators system
include("operators/operators.jl")

# Re-export main interface from comprehensive system
export create_operator, apply_operator!, laplace, gradient, divergence, curl
export poisson_solve, create_operator_set, OperatorSet
export DirichletBC, NeumannBC, RobinBC, PeriodicBC
export BoundaryConditionSet, demonstrate_operators

"""
    make_laplace_operator(grid::AbstractGrid) -> AbstractMatrix

Create discrete Laplacian operator matrix for the given grid.
"""
function make_laplace_operator(grid::AbstractCartesianGrid)
    # Implementation for Cartesian grids using finite differences
    total_points = prod(size(grid))
    dims = size(grid)
    
    # Create sparse Laplacian matrix
    I_vals = Int[]
    J_vals = Int[]
    V_vals = Float64[]
    
    for indices in CartesianIndices(dims)
        idx = LinearIndices(dims)[indices]
        
        # Diagonal term
        push!(I_vals, idx)
        push!(J_vals, idx) 
        push!(V_vals, -2 * ndim(grid))  # Will be scaled by 1/h²
        
        # Off-diagonal terms for each dimension
        for d in 1:ndim(grid)
            h = spacing(grid)[d]
            
            # Forward difference
            if indices[d] < dims[d]
                neighbor = CartesianIndex(ntuple(i -> i == d ? indices[i] + 1 : indices[i], ndim(grid)))
                neighbor_idx = LinearIndices(dims)[neighbor]
                push!(I_vals, idx)
                push!(J_vals, neighbor_idx)
                push!(V_vals, 1.0 / h^2)
            end
            
            # Backward difference  
            if indices[d] > 1
                neighbor = CartesianIndex(ntuple(i -> i == d ? indices[i] - 1 : indices[i], ndim(grid)))
                neighbor_idx = LinearIndices(dims)[neighbor]
                push!(I_vals, idx)
                push!(J_vals, neighbor_idx)
                push!(V_vals, 1.0 / h^2)
            end
        end
    end
    
    return sparse(I_vals, J_vals, V_vals, total_points, total_points)
end

function make_laplace_operator(grid::PolarSymGrid)
    # Laplacian in polar coordinates with symmetry: d²/dr² + (1/r)d/dr
    n = grid.num_r
    dr = spacing(grid)[1]
    
    I_vals = Int[]
    J_vals = Int[]
    V_vals = Float64[]
    
    for i in 1:n
        r = coordinate(grid, (i,))[1]
        
        if i == 1
            # Special handling at r=0 (center)
            push!(I_vals, i)
            push!(J_vals, i)
            push!(V_vals, -2.0 / dr^2)
            
            push!(I_vals, i)
            push!(J_vals, i+1)
            push!(V_vals, 2.0 / dr^2)
        elseif i == n
            # Boundary condition at r_max
            push!(I_vals, i)
            push!(J_vals, i)
            push!(V_vals, -2.0 / dr^2)
            
            push!(I_vals, i)
            push!(J_vals, i-1)
            push!(V_vals, 2.0 / dr^2)
        else
            # Interior points
            push!(I_vals, i)
            push!(J_vals, i)
            push!(V_vals, -2.0 / dr^2)
            
            # Forward difference
            push!(I_vals, i)
            push!(J_vals, i+1)
            push!(V_vals, 1.0 / dr^2 + 1.0 / (2 * r * dr))
            
            # Backward difference
            push!(I_vals, i)
            push!(J_vals, i-1)
            push!(V_vals, 1.0 / dr^2 - 1.0 / (2 * r * dr))
        end
    end
    
    return sparse(I_vals, J_vals, V_vals, n, n)
end

function make_laplace_operator(grid::SphericalSymGrid)
    # Laplacian in spherical coordinates with symmetry: d²/dr² + (2/r)d/dr
    n = grid.num_r
    dr = spacing(grid)[1]
    
    I_vals = Int[]
    J_vals = Int[]
    V_vals = Float64[]
    
    for i in 1:n
        r = coordinate(grid, (i,))[1]
        
        if i == 1
            # Special handling at r=0 (center)
            push!(I_vals, i)
            push!(J_vals, i)
            push!(V_vals, -6.0 / dr^2)
            
            push!(I_vals, i)
            push!(J_vals, i+1)
            push!(V_vals, 6.0 / dr^2)
        elseif i == n
            # Boundary condition at r_max
            push!(I_vals, i)
            push!(J_vals, i)
            push!(V_vals, -2.0 / dr^2)
            
            push!(I_vals, i)
            push!(J_vals, i-1)
            push!(V_vals, 2.0 / dr^2)
        else
            # Interior points
            push!(I_vals, i)
            push!(J_vals, i)
            push!(V_vals, -2.0 / dr^2)
            
            # Forward difference
            push!(I_vals, i)
            push!(J_vals, i+1)
            push!(V_vals, 1.0 / dr^2 + 1.0 / (r * dr))
            
            # Backward difference
            push!(I_vals, i)
            push!(J_vals, i-1)
            push!(V_vals, 1.0 / dr^2 - 1.0 / (r * dr))
        end
    end
    
    return sparse(I_vals, J_vals, V_vals, n, n)
end

"""
    make_gradient_operator(grid::AbstractGrid, dim::Int) -> AbstractMatrix

Create discrete gradient operator matrix for given dimension.
"""
function make_gradient_operator(grid::AbstractCartesianGrid, dim::Int)
    total_points = prod(size(grid))
    dims = size(grid)
    h = spacing(grid)[dim]
    
    I_vals = Int[]
    J_vals = Int[]
    V_vals = Float64[]
    
    for indices in CartesianIndices(dims)
        idx = LinearIndices(dims)[indices]
        
        if indices[dim] < dims[dim]
            # Forward difference
            neighbor = CartesianIndex(ntuple(i -> i == dim ? indices[i] + 1 : indices[i], ndim(grid)))
            neighbor_idx = LinearIndices(dims)[neighbor]
            
            push!(I_vals, idx)
            push!(J_vals, neighbor_idx)
            push!(V_vals, 1.0 / (2 * h))
        end
        
        if indices[dim] > 1
            # Backward difference
            neighbor = CartesianIndex(ntuple(i -> i == dim ? indices[i] - 1 : indices[i], ndim(grid)))
            neighbor_idx = LinearIndices(dims)[neighbor]
            
            push!(I_vals, idx)
            push!(J_vals, neighbor_idx)
            push!(V_vals, -1.0 / (2 * h))
        end
    end
    
    return sparse(I_vals, J_vals, V_vals, total_points, total_points)
end

"""
    make_divergence_operator(grid::AbstractGrid) -> AbstractMatrix

Create discrete divergence operator matrix.
"""
function make_divergence_operator(grid::AbstractCartesianGrid)
    # Divergence operator for vector fields
    # This would typically be more complex for vector field handling
    return make_gradient_operator(grid, 1)  # Simplified
end

"""
    make_curl_operator(grid::AbstractGrid) -> AbstractMatrix

Create discrete curl operator matrix.
"""
function make_curl_operator(grid::AbstractCartesianGrid)
    # Curl operator for vector fields
    # Implementation depends on dimensionality
    if ndim(grid) == 2
        # 2D curl: ∂v/∂x - ∂u/∂y
        grad_x = make_gradient_operator(grid, 1)
        grad_y = make_gradient_operator(grid, 2)
        # Would need proper vector field indexing
        return grad_x - grad_y  # Simplified
    else
        error("3D curl not implemented yet")
    end
end

# High-level operator functions that work with field data
"""
    grid_laplace(grid::AbstractGrid, field_data::AbstractArray) -> AbstractArray

Apply Laplacian operator to field data.
"""
function grid_laplace(grid::AbstractGrid, field_data::AbstractArray)
    L = make_laplace_operator(grid)
    return reshape(L * vec(field_data), size(field_data))
end

"""
    grid_gradient(grid::AbstractGrid, field_data::AbstractArray, dim::Int) -> AbstractArray

Apply gradient operator in given dimension to field data.
"""
function grid_gradient(grid::AbstractGrid, field_data::AbstractArray, dim::Int)
    G = make_gradient_operator(grid, dim)
    return reshape(G * vec(field_data), size(field_data))
end

"""
    grid_divergence(grid::AbstractGrid, vector_field_data::AbstractArray) -> AbstractArray

Apply divergence operator to vector field data.
"""
function grid_divergence(grid::AbstractGrid, vector_field_data::AbstractArray)
    D = make_divergence_operator(grid)
    # This would need proper handling of vector field structure
    return reshape(D * vec(vector_field_data[:,:,1]), size(vector_field_data)[1:end-1])
end

"""
    grid_curl(grid::AbstractGrid, vector_field_data::AbstractArray) -> AbstractArray

Apply curl operator to vector field data.
"""
function grid_curl(grid::AbstractGrid, vector_field_data::AbstractArray)
    C = make_curl_operator(grid)
    # This would need proper handling of vector field structure
    return reshape(C * vec(vector_field_data[:,:,1]), size(vector_field_data)[1:end-1])
end

# FFT-based operators for periodic boundaries
"""
    fft_laplace(grid::AbstractCartesianGrid, field_data::AbstractArray) -> AbstractArray

Fast Laplacian using FFT (for periodic boundaries).
"""
function fft_laplace(grid::AbstractCartesianGrid, field_data::AbstractArray)
    # Check if all dimensions are periodic
    bcs = get_boundary_conditions(grid)
    if !all(bc -> isa(bc, PeriodicBC), bcs)
        throw(ArgumentError("FFT Laplacian requires periodic boundary conditions"))
    end
    
    # Compute wavenumbers
    dims = size(grid)
    extents = extent(grid)
    
    # FFT of field
    field_fft = fft(field_data)
    
    # Apply Laplacian in Fourier space: -k²
    for indices in CartesianIndices(dims)
        k_squared = 0.0
        for d in 1:ndim(grid)
            L = last(extents[d]) - first(extents[d])
            k = 2π * (indices[d] - 1 - dims[d]÷2) / L
            k_squared += k^2
        end
        field_fft[indices] *= -k_squared
    end
    
    # Inverse FFT
    return real.(ifft(field_fft))
end

# =============================================================================
# BOUNDARY-AWARE DIFFERENTIAL OPERATORS
# =============================================================================

"""
    make_laplace_operator_with_bc(grid::AbstractGrid) -> AbstractMatrix

Create Laplacian operator that incorporates boundary conditions.
"""
function make_laplace_operator_with_bc(grid::AbstractGrid)
    # Start with basic operator
    L = make_laplace_operator(grid)
    
    # Modify for each dimension's boundary conditions
    bcs = get_boundary_conditions(grid)
    for (axis, bc) in enumerate(bcs)
        modify_operator_for_bc!(L, bc, grid, axis)
    end
    
    return L
end

"""
    grid_laplace_with_bc(grid::AbstractGrid, field_data::AbstractArray; time::Float64=0.0) -> AbstractArray

Apply Laplacian operator with boundary conditions enforced.
"""
function grid_laplace_with_bc(grid::AbstractGrid, field_data::AbstractArray; time::Float64=0.0)
    # Apply boundary conditions first
    field_with_bc = copy(field_data)
    apply_boundary_conditions!(field_with_bc, grid; time=time)
    
    # Apply Laplacian
    result = grid_laplace(grid, field_with_bc)
    
    # Apply boundary conditions to result if needed
    apply_boundary_conditions!(result, grid; time=time)
    
    return result
end

"""
    choose_laplace_method(grid::AbstractGrid) -> Symbol

Choose optimal Laplacian method based on boundary conditions and grid size.
"""
function choose_laplace_method(grid::AbstractGrid)
    bcs = get_boundary_conditions(grid)
    
    # Use FFT for all-periodic BCs
    if all(bc -> isa(bc, PeriodicBC), bcs)
        return :fft
    end
    
    # Use sparse matrix for small grids
    total_points = prod(size(grid))
    if total_points < 10000
        return :sparse
    end
    
    # Use finite difference for large grids with non-periodic BCs
    return :finite_difference
end

"""
    adaptive_laplace(grid::AbstractGrid, field_data::AbstractArray; time::Float64=0.0) -> AbstractArray

Apply Laplacian using the most appropriate method for the grid and boundary conditions.
"""
function adaptive_laplace(grid::AbstractGrid, field_data::AbstractArray; time::Float64=0.0)
    method = choose_laplace_method(grid)
    
    if method == :fft
        return fft_laplace(grid, field_data)
    elseif method == :sparse
        return grid_laplace_with_bc(grid, field_data; time=time)
    else  # :finite_difference
        return finite_difference_laplace_with_bc(grid, field_data; time=time)
    end
end

"""
    finite_difference_laplace_with_bc(grid::AbstractGrid, field_data::AbstractArray; time::Float64=0.0) -> AbstractArray

Direct finite difference Laplacian with boundary condition handling.
"""
function finite_difference_laplace_with_bc(grid::AbstractGrid, field_data::AbstractArray; time::Float64=0.0)
    # Apply boundary conditions
    field_with_bc = copy(field_data)
    apply_boundary_conditions!(field_with_bc, grid; time=time)
    
    # Initialize result
    result = similar(field_data)
    dims = size(grid)
    spacings = spacing(grid)
    
    # Apply finite difference stencil with BC-aware indexing
    for indices in CartesianIndices(dims)
        laplacian_val = 0.0
        
        for d in 1:ndim(grid)
            h = spacings[d]
            
            # Get neighbor indices with boundary handling
            idx_low = get_bc_aware_index(indices, d, -1, dims, grid, d)
            idx_high = get_bc_aware_index(indices, d, +1, dims, grid, d)
            
            # Second-order finite difference
            val_center = field_with_bc[indices]
            val_low = field_with_bc[idx_low]
            val_high = field_with_bc[idx_high]
            
            laplacian_val += (val_low - 2*val_center + val_high) / h^2
        end
        
        result[indices] = laplacian_val
    end
    
    return result
end

"""
    get_bc_aware_index(center_idx, axis, offset, dims, grid, bc_axis)

Get index for finite difference stencil with boundary condition awareness.
"""
function get_bc_aware_index(center_idx::CartesianIndex, axis::Int, offset::Int, dims::Tuple, grid::AbstractGrid, bc_axis::Int)
    new_coords = [center_idx[i] for i in 1:length(dims)]
    new_coords[axis] += offset
    
    # Handle boundary conditions
    if new_coords[axis] < 1
        bc = get_boundary_condition(grid, bc_axis)
        if isa(bc, PeriodicBC)
            new_coords[axis] = dims[axis] + new_coords[axis]  # Wrap around
        else
            # For other BCs, ghost cells are handled by apply_boundary_conditions!
            new_coords[axis] = 1  # Use boundary value
        end
    elseif new_coords[axis] > dims[axis]
        bc = get_boundary_condition(grid, bc_axis)
        if isa(bc, PeriodicBC)
            new_coords[axis] = new_coords[axis] - dims[axis]  # Wrap around
        else
            new_coords[axis] = dims[axis]  # Use boundary value
        end
    end
    
    return CartesianIndex(new_coords...)
end

# =============================================================================
# INTEGRATION HELPER FUNCTIONS
# =============================================================================

"""
    migrate_to_comprehensive_operators()

Print information about migrating to the new comprehensive operator system.
"""
function migrate_to_comprehensive_operators()
    println("""
    PDEJulia.jl Operator System Migration Guide:
    ==========================================
    
    The PDEJulia.jl library now includes a comprehensive operator system
    with enhanced functionality. Consider migrating to the new interface:
    
    Old interface:              New interface:
    grid_laplace(grid, field)   laplace(field, grid, bcs)
    make_laplace_operator(grid) create_operator(:laplacian, grid, bcs)
    
    New features available:
    • Comprehensive boundary condition support
    • Operator caching and optimization
    • All coordinate systems (Cartesian, spherical, cylindrical, polar)
    • High-order finite differences
    • FFT-based operators
    • Performance benchmarking
    • Automatic operator selection
    
    Example usage:
    ```julia
    # Create boundary conditions
    bcs = create_bc_set_cartesian_2d(
        DirichletBC(0.0), DirichletBC(0.0),
        DirichletBC(0.0), DirichletBC(0.0)
    )
    
    # Use high-level interface
    result = laplace(field, grid, bcs)
    
    # Or create operator for repeated use
    laplacian = create_operator(:laplacian, grid, bcs)
    laplacian(field, result)
    ```
    """)
end